import torch
import torchvision
import torch.utils.data
import os
import pandas
import json
import xml.dom.minidom
import PIL.Image


class COCO(torch.utils.data.Dataset):
    def __init__(self, root, image_dir, annotations_file, transform=None, target_transform=None):
        self.image_path = os.path.join(root, image_dir)
        self.transform = transform
        with open(os.path.join(root, annotations_file), 'r') as f:
            self.labels_json = json.load(f)
        self.target_transform = target_transform
        self.ids = []
        for img in self.labels_json['images']:
            self.ids.append([])
        for i, ann in enumerate(self.labels_json['annotations']):
            self.ids[i].append(ann['id'])

    def __len__(self):
        return len(self.labels_json['images'])

    def __getitem__(self, idx):
        image = PIL.Image.open(os.path.join(self.image_path, self.labels_json['images'][idx]['file_name']))
        image_size = image.size
        ann_ids = self.ids[idx]
        ann_s = []
        for ann_id in ann_ids:
            ann = self.labels_json['annotations'][ann_id]
            category_id = ann['category_id']
            bbox = ann['bbox']
            bbox[0] = bbox[0] / image_size[0]
            bbox[1] = bbox[1] / image_size[1]
            bbox[2] = bbox[2] / image_size[0]
            bbox[3] = bbox[3] / image_size[1]
            bbox.append(float(category_id))
            label = torch.tensor(bbox)

            ann_s.append(label)
        if self.transform:
            image = self.transform(image)

        if self.target_transform:
            for i, l in enumerate(ann_s):
                ann_s[i] = self.target_transform(l)
        return image, torch.stack(ann_s)


class DogDetectionDataset(torch.utils.data.Dataset):
    def __init__(self, root, image_dir, annotations_file, transform=None, target_transform=None):
        self.image_path = os.path.join(root, image_dir)
        self.transform = transform
        self.img_labels = pandas.read_csv(os.path.join(root, annotations_file))
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        image = torchvision.io.read_image(os.path.join(self.image_path, self.img_labels.iloc[idx][0])).to(torch.float32)

        image_scale = torch.tensor([image.shape[2], image.shape[1], image.shape[2], image.shape[1]],
                                   dtype=torch.float32)

        label = torch.tensor(self.img_labels.iloc[idx][1:], dtype=torch.float32)
        label[1:] = label[1:] / image_scale
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label


class dog_dataset(torch.utils.data.Dataset):
    def __init__(self, root, image, label, transform=None, target_transform=None):
        self.img_dir = os.path.join(root, image)
        self.xml_dir = os.path.join(root, label)
        dogs_xml_tag = []

        for path in os.listdir(self.xml_dir):
            current_xml = os.path.join(self.xml_dir, path)
            dog_xml_tag = readxml(current_xml)
            dogs_xml_tag.append(torch.Tensor(dog_xml_tag))
        # print((dogs_xml_tag))
        all_img_path = []
        for path in os.listdir(self.img_dir):
            current_img = os.path.join(self.img_dir, path)
            all_img_path.append(current_img)

        self.img_labels = dogs_xml_tag
        self.img_path = all_img_path
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        image = torchvision.io.read_image(self.img_path[idx]).to(torch.float32)
        image_scale = torch.tensor([image.shape[2], image.shape[1], image.shape[2], image.shape[1]],
                                   dtype=torch.float32)
        label_box = self.img_labels[idx] / image_scale
        label = torch.cat([torch.zeros(1), label_box])
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label


def readxml(filename):
    dom_tree = xml.dom.minidom.parse(filename)
    annotation = dom_tree.documentElement
    folder_name = annotation.getElementsByTagName('folder')[0].childNodes[0].data
    image_filename = annotation.getElementsByTagName('filename')[0].childNodes[0].data
    class_name = annotation.getElementsByTagName('object')[0].childNodes[1].childNodes[0].data
    objs = annotation.getElementsByTagName('object')
    for obj in objs:
        bndboxes = obj.getElementsByTagName('bndbox')[0]
        name = obj.getElementsByTagName('name')[0].childNodes[0].data
        if name != 'dog':
            continue
        xmin = bndboxes.getElementsByTagName('xmin')[0].childNodes[0].data
        ymin = bndboxes.getElementsByTagName('ymin')[0].childNodes[0].data
        xmax = bndboxes.getElementsByTagName('xmax')[0].childNodes[0].data
        ymax = bndboxes.getElementsByTagName('ymax')[0].childNodes[0].data

        return [int(xmin), int(ymin), int(xmax), int(ymax)]
